from configure import get_args
from dataLoader import DataLoader, load_data
from crnnLSTM import buildCrnn
from visualization import resultTest
import numpy as np

# based on baseline datasets

if __name__ == '__main__':
    args = get_args(418, 53)
    X_train,Y_train = load_data(args,args.train_root_path, args.AD_dir, args.CN_dir, True)
    X_val, Y_val = load_data(args, args.val_root_path, args.AD_dir, args.CN_dir, False)
    X_test, Y_test = load_data(args, args.test_root_path, args.AD_dir , args.CN_dir, False)
    crnn_model = buildCrnn(args)
    from sklearn.utils.class_weight import compute_class_weight
    from tensorflow.keras.callbacks import ReduceLROnPlateau, EarlyStopping
    from keras.callbacks import ModelCheckpoint
    Y_train_labels = np.argmax(Y_train, axis=1)
    class_weights = compute_class_weight('balanced', classes=np.unique(Y_train_labels), y=Y_train_labels)
    class_weights_dict = dict(enumerate(class_weights))
    checkpoint_callback = ModelCheckpoint(
    filepath= args.model_saving_path, 
    monitor=args.monitor,        
    save_best_only=True,        
    mode=args.mode,                 
    verbose=args.verbose                   
    )
    reduce_lr = ReduceLROnPlateau(monitor=args.monitor, factor=args.factor, patience=args.lr_patience, min_lr=args.min_lr, verbose=args.verbose)
    early_stopping = EarlyStopping(monitor=args.early_stopping_monitor, patience=args.early_stopping_patience, restore_best_weights=True)
    history = crnn_model.fit(X_train, 
                             validation_data=(X_val, Y_val), 
                             epochs=args.epochs, 
                             steps_per_epoch=args.steps_per_epoch,
                             callbacks=[reduce_lr, early_stopping, checkpoint_callback],
                             class_weight=class_weights_dict)
    
    resultTest(history, crnn_model, X_test, Y_test)

    
    
    
    
